# python <pdb_name> <ori_chains> <aim_chains>
# python example.pdb ABC EFG

import sys

def chainchange(pdb_name,ori_chains,to_chains ):
    ori_chains_list = list(ori_chains)
    to_chains_list = list(to_chains)

    if len(ori_chains_list) != len(to_chains):
        print("输入的’初始链名‘数量应该等于’结果链名‘")
        sys.exit(0)

    pdb = open("./{}.pdb".format(pdb_name.split(".pdb")[0]), "r")
    pdb_lines = [i for i in pdb]

    pdb.close()
    for i in range(len(pdb_lines)):
        if pdb_lines[i][0:4] == 'ATOM':
            for j in range(len(ori_chains_list)):
                if pdb_lines[i][21] == ori_chains_list[j]:
                    pdb_lines[i] = pdb_lines[i][0:21] + to_chains_list[j] + pdb_lines[i][22:-1] + pdb_lines[i][-1]

    pdb = open("./{}_chain_changed.pdb".format(pdb_name), "w")
    pdb.writelines(pdb_lines)
    pdb.close()